#!/usr/bin/env python
################################################################
#
# force2imd:
#   convert a potfit configuration (force file) into
#   an IMD checkpoint file
#
################################################################
#
#   Copyright 2013
#             Institute for Theoretical and Applied Physics
#             University of Stuttgart, D-70550 Stuttgart, Germany
#             https://www.potfit.net/
#
#################################################################
#
#   This file is part of potfit.
#
#   potfit is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   potfit is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with potfit; if not, see <http://www.gnu.org/licenses/>.
#
#################################################################

import argparse
import os
import sys
import toe

parser = argparse.ArgumentParser(description='Convert a potfit configuration into an IMD checkpoint file.')
parser.add_argument('-n', '--number', type=int, default=0, help='number of the configuration to convert')
parser.add_argument('filename', type=str, help='name of the potfit config file')
args = parser.parse_args()

# open configuration file
try:
    f = open(args.filename)
except IOError as e:
    sys.stderr.write("Could not open the configuration file:\n")
    sys.stderr.write("%s\n" % e)
    sys.exit()

number = args.number

# fast forward to requested configuration
if number > 0:
    f.readline()
    while number > 0:
        temp = f.readline()
        if temp[0] == '#' and temp[1] == 'N':
            number = number - 1
else:
    temp = f.readline()

# read header
while temp[0] == '#' and temp[1] != 'F':
    if temp[1] == 'N':
        natoms = temp.strip().split()[1]
    if temp[1] == 'C':
        elements = temp.strip().split()[1:]
    if temp[1] == 'X':
        vect_x = [float(s) for s in temp.strip().split()[1:]]
    if temp[1] == 'Y':
        vect_y = [float(s) for s in temp.strip().split()[1:]]
    if temp[1] == 'Z':
        vect_z = [float(s) for s in temp.strip().split()[1:]]
    temp = f.readline()

try:
    if elements != []:
        pass
except:
    sys.stderr.write("ERROR: \n")
    sys.stderr.write("  The #C line was not found for the {}. configuration in {}!\n".format(args.number,args.filename))
    sys.exit()

# write header
print "#F A 1 1 1 3 0 0"
print "#C number type mass x y z"
print "#X {:12.8f} {:12.8f} {:12.8f}".format(vect_x[0],vect_x[1],vect_x[2])
print "#Y {:12.8f} {:12.8f} {:12.8f}".format(vect_y[0],vect_y[1],vect_y[2])
print "#Z {:12.8f} {:12.8f} {:12.8f}".format(vect_z[0],vect_z[1],vect_z[2])
print "## Generated by {} from {} (config {})".format(sys.argv[0],os.path.abspath(args.filename),args.number)
print "#E"

# write atom data
temp = f.readline()
i = 0
while temp != '' and temp[0] != '#':
    data = temp.strip().split()
    atom_type = int(data[0])
    mass = toe.getMassByName(elements[atom_type])
    if mass == 0.0:
        mass = data[0] + 1
    position = [float(s) for s in data[1:4]]
    print "{:6}\t{}\t{}\t{:10.6f} {:10.6f} {:10.6f}".format(i, data[0],mass,position[0],position[1],position[2])
    i = i + 1
    temp = f.readline()
